'''
Created on Aug 31, 2009

@author: mkiyer
'''
       
from veggie.sample.sampledb2 import get_sampledb
import xml.etree.cElementTree as etree
import logging

def validate_samples(samples):
    validated = []
    for s in samples:
        if s in get_sampledb().param_schema['oncoseq_sample_name']:
            validated.append(s)
        else:
            logging.error('sample: %s not found in sampledb' % s)
            raise KeyError('sample %s not found in sampledb' % s)
    return validated

def validate_cond(cond):        
    for k, v in cond.iteritems():
        if k not in get_sampledb().param_schema:
            print 'cond: key', k, 'not found in', get_sampledb().param_schema.keys()
            return False
        else:
            if v not in get_sampledb().param_schema[k]:
                print 'cond: key', k, 'value', v, 'not found in', get_sampledb().param_schema[k]
                return False
    return True

def find_samples_matching_cond(cond):
    if cond is None:
        return []
    if not validate_cond(cond):
        return []
    matching_samples = []
    for sample_name, sample_params in get_sampledb().params.iteritems():
        props = sample_params
        match = True
        for k, v in cond.iteritems():                
            if k not in props or props[k] != v:
                match = False
                break
        if match == True:
            matching_samples.append(sample_name)
    return matching_samples

def parse_elem_samplegroup(elem):
    # parse the xml elements
    assert elem.tag == 'samplegroup'
    name_elem = elem.find('name')
    name = name_elem.text if name_elem is not None else None
    color_elem = elem.find('color')
    color = color_elem.text if color_elem is not None else None
    cond_elems = elem.findall('cond')
    conds = []
    if cond_elems is not None:
        for cond_elem in cond_elems:
            cond = {}
            for e in cond_elem.getchildren():
                cond[e.tag] = e.text
            conds.append(cond)    
    samples = [x.text for x in elem.findall('sample')]
    # now build a samplegroup object from the elements
    return SampleGroup(name, color, conds, samples)

def parse_samplegroups_xml(xmlfile):
    tree = etree.parse(xmlfile)
    root = tree.getroot()
    assert root.tag == 'samplegroups'    
    for e in root.getchildren():
        yield parse_elem_samplegroup(e)
    
class SampleGroup(object):
    '''
    Represents a group of samples and associates each group with a name
    and color for plotting
    
    A SampleGroup is specified by either a list of sample names or a 
    list of conditions (conds) that use SampleDB parameters to describe a
    class of samples
    '''
    def __init__(self, name=None, color='black', conds=None, samples=None):
        self.name = name
        self.color = color

        self.samples = set([])        
        if samples is not None:
            self.samples.update(validate_samples(samples))

        self.conds = []        
        if conds is not None:         
            for cond in conds:
                if validate_cond(cond):
                    self.conds.append(cond)
                    self.samples.update(find_samples_matching_cond(cond))


if __name__ == '__main__':
    import sys
    groups = [x for x in parse_samplegroups_xml(sys.argv[1])]
    